{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "443c48f0",
   "metadata": {},
   "source": [
    "# CondenseNet 关键观察与创新点总结\n",
    "\n",
    "**CondenseNet** 是一种高效的卷积神经网络架构，结合了 **密集连接（DenseNet）** 和 **可学习分组卷积（Learnable Group Convolution, LGC）** 的思想，旨在实现高性能和低计算成本的图像分类任务。以下是其主要的关键观察与创新点。\n",
    "\n",
    "---\n",
    "\n",
    "## 🔍 关键观察（Key Observations）\n",
    "\n",
    "1. **特征复用是高效模型设计的核心**\n",
    "   - DenseNet 中通过密集连接实现了特征复用，减少了冗余信息的学习，提升了模型效率。\n",
    "   - CondenseNet 延续了这一理念，并进一步优化了特征通道的选择机制。\n",
    "\n",
    "2. **冗余通道对性能提升有限**\n",
    "   - 实验发现，许多通道在推理过程中贡献较小，去除这些“非必要”通道对精度影响极小，却能显著降低计算量。\n",
    "\n",
    "3. **训练阶段压缩通道可实现推理加速**\n",
    "   - 在训练过程中逐步淘汰不重要的通道（即“condense”），使网络在推理时仅保留必要通道，从而减少FLOPs和参数量。\n",
    "\n",
    "4. **结构稀疏性优于随机稀疏性**\n",
    "   - 相比于随机剪枝或正则化方法，CondenseNet 采用结构化的通道剪枝方式，更易于部署和加速。\n",
    "\n",
    "---\n",
    "\n",
    "## 🧩 创新点（Innovations）\n",
    "\n",
    "### 1. ✅ 可学习分组卷积（Learnable Group Convolution, LGC）\n",
    "\n",
    "- 在训练过程中，将标准卷积分解为多个组卷积。\n",
    "- 每个输入通道只连接到一组输出通道，且这种连接关系在训练中通过反向传播自动学习。\n",
    "- 推理时，只保留每组中响应最强的输出通道，实现通道压缩。\n",
    "\n",
    "### 2. 🚀 分阶段训练策略（Progressive Training）\n",
    "\n",
    "- 将训练分为多个阶段，在每个阶段逐渐增加压缩程度。\n",
    "- 类似课程学习的方式，让网络逐步适应稀疏结构，避免一次性压缩带来的性能下降。\n",
    "\n",
    "  <img src=\"resources/condensenet_condensing_procedure.png\" alt=\"drawing\" width=\"60%\"/>\n",
    "\n",
    "### 3. 📉 结构化稀疏 + 高效推理\n",
    "\n",
    "- 所有压缩操作都是结构化的，可以在通用硬件（如GPU）上直接加速，无需特殊硬件支持。\n",
    "- 模型大小和计算量显著降低，同时保持了较高的准确率。\n",
    "\n",
    "### 4. 🔁 灵活适配不同压缩率\n",
    "\n",
    "- 可以根据实际需求调整压缩阶段数和每组的通道数，灵活平衡精度与速度。\n",
    "\n",
    "---\n",
    "\n",
    "<img src=\"resources/condensenet_blocks.png\" alt=\"drawing\" width=\"50%\"/>\n",
    "\n",
    "## ✅ 总结\n",
    "\n",
    "CondenseNet 提出了一种新颖的、结构化稀疏化方法，结合了 DenseNet 的高效特征复用和可学习分组卷积的思想。其核心在于通过训练过程中的“压缩”机制，引导网络在推理阶段仅保留必要的通道连接，从而实现高精度与高效率的统一。\n",
    "\n",
    "如果你正在寻找一个 **轻量级但准确率高的网络架构**，CondenseNet 是一个非常值得尝试的选择！\n",
    "\n",
    "> 💡 **提示**：如果你需要我继续补充代码示例、训练技巧或 PyTorch 实现要点，也可以告诉我！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b93e5450",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Device:  cuda\n"
     ]
    }
   ],
   "source": [
    "# 自动重新加载外部module，使得修改代码之后无需重新import\n",
    "# see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from hdd.device.utils import get_device\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# 设置训练数据的路径\n",
    "DATA_ROOT = \"~/workspace/hands-dirty-on-dl/dataset\"\n",
    "# 设置TensorBoard的路径\n",
    "TENSORBOARD_ROOT = \"~/workspace/hands-dirty-on-dl/dataset\"\n",
    "# 设置预训练模型参数路径\n",
    "TORCH_HUB_PATH = \"~/workspace/hands-dirty-on-dl/pretrained_models\"\n",
    "torch.hub.set_dir(TORCH_HUB_PATH)\n",
    "# 挑选最合适的训练设备\n",
    "DEVICE = get_device([\"cuda\", \"cpu\"])\n",
    "print(\"Device: \", DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6746460f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "from hdd.data_util.auto_augmentation import CIFAR10Policy\n",
    "\n",
    "# 训练超参数和数据增强来自 https://github.com/omihub777/ViT-CIFAR\n",
    "CIFAR_10_MEAN = [0.4914, 0.4822, 0.4465]\n",
    "CIFAR_10_STD = [0.2470, 0.2435, 0.2616]\n",
    "BATCH_SIZE = 128\n",
    "\n",
    "val_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(CIFAR_10_MEAN, CIFAR_10_STD),\n",
    "    ]\n",
    ")\n",
    "\n",
    "val_dataloader = torch.utils.data.DataLoader(\n",
    "    datasets.CIFAR10(\n",
    "        root=DATA_ROOT, train=False, download=True, transform=val_transform\n",
    "    ),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=False,\n",
    "    num_workers=8,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "train_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.RandomCrop(size=32, padding=4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        CIFAR10Policy(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(CIFAR_10_MEAN, CIFAR_10_STD),\n",
    "    ]\n",
    ")\n",
    "\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    datasets.CIFAR10(\n",
    "        root=DATA_ROOT, train=True, download=True, transform=train_transform\n",
    "    ),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    num_workers=8,\n",
    "    pin_memory=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "51cf80cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from dataclasses import dataclass\n",
    "from typing import Callable, List, Optional, Tuple\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data.dataset import Dataset\n",
    "from hdd.train.classification_utils import _eval_classifier_naive\n",
    "from hdd.train.early_stopping import EarlyStoppingInterface\n",
    "\n",
    "\n",
    "def train_classifier_one_epoch(\n",
    "    net: nn.Module,\n",
    "    criteria: nn.CrossEntropyLoss,\n",
    "    optimizer: optim.Optimizer,\n",
    "    train_loader: torch.utils.data.DataLoader,\n",
    "    device: torch.device,\n",
    "    epoch: int,\n",
    "    max_epochs: int,\n",
    "    lasso_lambda: float,\n",
    ") -> Tuple[float, float]:\n",
    "    \"\"\"Naive training procedure to train classifier for one epoch.\n",
    "\n",
    "    Args:\n",
    "        net: network instance.\n",
    "        criteria: Loss function. Typically nn.CrossEntropyLoss\n",
    "        optimizer: optimizer.\n",
    "        train_loader: train data\n",
    "        device: device to run the training.\n",
    "\n",
    "    Returns:\n",
    "        avg train loss and train accuracy.\n",
    "    \"\"\"\n",
    "\n",
    "    train_loss = 0.0\n",
    "    correct_items = 0\n",
    "    total_items = 0\n",
    "    net.train()\n",
    "    progress = epoch / max_epochs\n",
    "    learned_module_list = []\n",
    "    for m in net.modules():\n",
    "        if m.__str__().startswith(\"LearnedGroupConv\"):\n",
    "            learned_module_list.append(m)\n",
    "\n",
    "    for i, [Xs, ys] in enumerate(train_loader):\n",
    "        Xs, ys = Xs.to(device), ys.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        progress = float(epoch * len(train_loader) + i) / (\n",
    "            max_epochs * len(train_loader)\n",
    "        )\n",
    "        logits = net(Xs, progress)\n",
    "        loss = criteria(logits, ys)\n",
    "        lasso_loss = 0\n",
    "        for m in learned_module_list:\n",
    "            lasso_loss = lasso_loss + m.lasso_loss\n",
    "            loss = loss + lasso_lambda * lasso_loss\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "        correct_items += torch.sum(torch.argmax(logits, dim=1) == ys).item()\n",
    "        total_items += Xs.shape[0]\n",
    "\n",
    "    avg_train_loss = train_loss / len(train_loader)\n",
    "    accuracy = correct_items / total_items\n",
    "    return avg_train_loss, accuracy\n",
    "\n",
    "\n",
    "def train_condensenet_model(\n",
    "    net: nn.Module,\n",
    "    criteria,\n",
    "    max_epochs: int,\n",
    "    lasso_lambda: float,\n",
    "    train_loader: torch.utils.data.DataLoader,\n",
    "    val_loader: torch.utils.data.DataLoader,\n",
    "    device: torch.device,\n",
    "    optimizer: optim.Optimizer,\n",
    "    scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,\n",
    "    early_stopper: Optional[EarlyStoppingInterface] = None,\n",
    "    verbose: bool = True,\n",
    "    eval_classifier: Callable[\n",
    "        [\n",
    "            nn.Module,\n",
    "            nn.CrossEntropyLoss,\n",
    "            torch.utils.data.DataLoader,\n",
    "            torch.device,\n",
    "        ],\n",
    "        Tuple[float, float],\n",
    "    ] = _eval_classifier_naive,\n",
    ") -> dict[str, list[float]]:\n",
    "    \"\"\"Naive classifier training procedure.\n",
    "\n",
    "    Args:\n",
    "        net: classification model.\n",
    "        criteria: loss function.\n",
    "        max_epochs: maximum number of epochs.\n",
    "        train_loader: train dataloader.\n",
    "        val_loader: validation dataloader.\n",
    "        device: network device.\n",
    "        optimizer: optimizer\n",
    "        scheduler: learning rate scheduler.\n",
    "        early_stopper: early stopper.\n",
    "        verbose: Print anything or not. Defaults to True.\n",
    "        eval_classifier: Function to eval the classifier for one epoch.\n",
    "    Returns:\n",
    "        training statistics.\n",
    "    \"\"\"\n",
    "    result = {\n",
    "        \"train_loss\": [],\n",
    "        \"val_loss\": [],\n",
    "        \"train_accuracy\": [],\n",
    "        \"val_accuracy\": [],\n",
    "    }\n",
    "    for epoch in range(1, max_epochs + 1):\n",
    "        t0 = time.time()\n",
    "        avg_train_loss, train_accuracy = train_classifier_one_epoch(\n",
    "            net,\n",
    "            criteria,\n",
    "            optimizer,\n",
    "            train_loader,\n",
    "            device,\n",
    "            epoch,\n",
    "            max_epochs,\n",
    "            lasso_lambda,\n",
    "        )\n",
    "        t1 = time.time()\n",
    "        if scheduler is not None:\n",
    "            scheduler.step()\n",
    "        avg_val_loss, val_accuracy = eval_classifier(\n",
    "            net,\n",
    "            criteria,\n",
    "            val_loader,\n",
    "            device,\n",
    "        )\n",
    "        if verbose:\n",
    "            print(\n",
    "                f\"Epoch: {epoch}/{max_epochs} \"\n",
    "                f\"Train Loss: {avg_train_loss:0.4f} \"\n",
    "                f\"Accuracy: {train_accuracy:0.4f} \"\n",
    "                f\"Time: {t1 - t0:0.5f} \"\n",
    "                f\" | Val Loss: {avg_val_loss:0.4f} \"\n",
    "                f\"Accuracy: {val_accuracy:0.4f}\"\n",
    "            )\n",
    "        result[\"train_loss\"].append(avg_train_loss)\n",
    "        result[\"val_loss\"].append(avg_val_loss)\n",
    "        result[\"train_accuracy\"].append(train_accuracy)\n",
    "        result[\"val_accuracy\"].append(val_accuracy)\n",
    "        if early_stopper is not None:\n",
    "            if early_stopper(val_loss=avg_val_loss, model=net):\n",
    "                print(f\"Early stop at epoch {epoch}!\")\n",
    "                early_stopper.load_best_model(net)\n",
    "                return result\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "45ef1441",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#Parameter: 1451594\n",
      "Epoch: 1/100 Train Loss: 6.5175 Accuracy: 0.0970 Time: 20.42050  | Val Loss: 2.3529 Accuracy: 0.0939\n",
      "Epoch: 2/100 Train Loss: 6.1064 Accuracy: 0.2585 Time: 19.66456  | Val Loss: 1.8150 Accuracy: 0.3684\n",
      "Epoch: 3/100 Train Loss: 5.5055 Accuracy: 0.3847 Time: 20.18755  | Val Loss: 1.5877 Accuracy: 0.5052\n",
      "Epoch: 4/100 Train Loss: 4.9899 Accuracy: 0.4660 Time: 19.55599  | Val Loss: 1.4525 Accuracy: 0.5785\n",
      "Epoch: 5/100 Train Loss: 4.5329 Accuracy: 0.5312 Time: 19.74438  | Val Loss: 1.3364 Accuracy: 0.6381\n",
      "Epoch: 6/100 Train Loss: 4.1239 Accuracy: 0.5744 Time: 19.97007  | Val Loss: 1.2596 Accuracy: 0.6717\n",
      "Epoch: 7/100 Train Loss: 3.7005 Accuracy: 0.6214 Time: 19.78455  | Val Loss: 1.2155 Accuracy: 0.6907\n",
      "Epoch: 8/100 Train Loss: 3.3005 Accuracy: 0.6590 Time: 19.96895  | Val Loss: 1.1969 Accuracy: 0.7079\n",
      "Epoch: 9/100 Train Loss: 2.9242 Accuracy: 0.6899 Time: 19.68217  | Val Loss: 1.0810 Accuracy: 0.7526\n",
      "Epoch: 10/100 Train Loss: 2.5992 Accuracy: 0.7066 Time: 19.76372  | Val Loss: 1.0602 Accuracy: 0.7678\n",
      "Epoch: 11/100 Train Loss: 2.3243 Accuracy: 0.7146 Time: 19.91731  | Val Loss: 1.0292 Accuracy: 0.7744\n",
      "Epoch: 12/100 Train Loss: 2.0936 Accuracy: 0.7331 Time: 19.53302  | Val Loss: 1.0577 Accuracy: 0.7733\n",
      "Epoch: 13/100 Train Loss: 1.9287 Accuracy: 0.7427 Time: 19.42895  | Val Loss: 0.9960 Accuracy: 0.7961\n",
      "Epoch: 14/100 Train Loss: 1.8117 Accuracy: 0.7518 Time: 19.46499  | Val Loss: 0.9566 Accuracy: 0.8117\n",
      "Epoch: 15/100 Train Loss: 1.7193 Accuracy: 0.7607 Time: 19.79444  | Val Loss: 0.9699 Accuracy: 0.8015\n",
      "Epoch: 16/100 Train Loss: 1.6316 Accuracy: 0.7692 Time: 19.93855  | Val Loss: 0.9024 Accuracy: 0.8296\n",
      "Epoch: 17/100 Train Loss: 1.5491 Accuracy: 0.7737 Time: 20.30779  | Val Loss: 0.9207 Accuracy: 0.8229\n",
      "Epoch: 18/100 Train Loss: 1.5024 Accuracy: 0.7809 Time: 19.94830  | Val Loss: 0.9380 Accuracy: 0.8220\n",
      "Epoch: 19/100 Train Loss: 1.4656 Accuracy: 0.7880 Time: 19.30747  | Val Loss: 0.8776 Accuracy: 0.8417\n",
      "Epoch: 20/100 Train Loss: 1.4297 Accuracy: 0.7922 Time: 18.76424  | Val Loss: 0.8871 Accuracy: 0.8350\n",
      "Epoch: 21/100 Train Loss: 1.4001 Accuracy: 0.7966 Time: 19.18311  | Val Loss: 0.8655 Accuracy: 0.8439\n",
      "Epoch: 22/100 Train Loss: 1.3786 Accuracy: 0.7999 Time: 20.10188  | Val Loss: 0.8600 Accuracy: 0.8480\n",
      "Epoch: 23/100 Train Loss: 1.3548 Accuracy: 0.8068 Time: 19.40838  | Val Loss: 0.8340 Accuracy: 0.8576\n",
      "Epoch: 24/100 Train Loss: 1.3358 Accuracy: 0.8119 Time: 20.06569  | Val Loss: 0.9165 Accuracy: 0.8300\n",
      "Epoch: 25/100 Train Loss: 1.3091 Accuracy: 0.8197 Time: 20.12450  | Val Loss: 0.8231 Accuracy: 0.8640\n",
      "Epoch: 26/100 Train Loss: 1.3003 Accuracy: 0.8153 Time: 20.01105  | Val Loss: 0.8211 Accuracy: 0.8642\n",
      "Epoch: 27/100 Train Loss: 1.2804 Accuracy: 0.8203 Time: 19.95559  | Val Loss: 0.8414 Accuracy: 0.8582\n",
      "Epoch: 28/100 Train Loss: 1.2700 Accuracy: 0.8218 Time: 20.07592  | Val Loss: 0.8122 Accuracy: 0.8731\n",
      "Epoch: 29/100 Train Loss: 1.2551 Accuracy: 0.8247 Time: 20.43467  | Val Loss: 0.8139 Accuracy: 0.8698\n",
      "Epoch: 30/100 Train Loss: 1.2432 Accuracy: 0.8266 Time: 19.97511  | Val Loss: 0.8099 Accuracy: 0.8715\n",
      "Epoch: 31/100 Train Loss: 1.2291 Accuracy: 0.8299 Time: 20.73610  | Val Loss: 0.8222 Accuracy: 0.8694\n",
      "Epoch: 32/100 Train Loss: 1.2146 Accuracy: 0.8355 Time: 19.67448  | Val Loss: 0.7866 Accuracy: 0.8779\n",
      "Epoch: 33/100 Train Loss: 1.1923 Accuracy: 0.8358 Time: 20.10113  | Val Loss: 0.8023 Accuracy: 0.8725\n",
      "Epoch: 34/100 Train Loss: 1.1738 Accuracy: 0.8386 Time: 19.55709  | Val Loss: 0.7834 Accuracy: 0.8778\n",
      "Epoch: 35/100 Train Loss: 1.1663 Accuracy: 0.8414 Time: 19.49982  | Val Loss: 0.7877 Accuracy: 0.8793\n",
      "Epoch: 36/100 Train Loss: 1.1540 Accuracy: 0.8451 Time: 20.01287  | Val Loss: 0.7969 Accuracy: 0.8724\n",
      "Epoch: 37/100 Train Loss: 1.1446 Accuracy: 0.8466 Time: 19.63873  | Val Loss: 0.8082 Accuracy: 0.8721\n",
      "Epoch: 38/100 Train Loss: 1.1396 Accuracy: 0.8481 Time: 19.66327  | Val Loss: 0.7924 Accuracy: 0.8764\n",
      "Epoch: 39/100 Train Loss: 1.1321 Accuracy: 0.8502 Time: 19.55413  | Val Loss: 0.7657 Accuracy: 0.8897\n",
      "Epoch: 40/100 Train Loss: 1.1295 Accuracy: 0.8485 Time: 19.13028  | Val Loss: 0.7690 Accuracy: 0.8869\n",
      "Epoch: 41/100 Train Loss: 1.1129 Accuracy: 0.8545 Time: 19.48546  | Val Loss: 0.7578 Accuracy: 0.8916\n",
      "Epoch: 42/100 Train Loss: 1.1091 Accuracy: 0.8557 Time: 19.27842  | Val Loss: 0.7475 Accuracy: 0.8946\n",
      "Epoch: 43/100 Train Loss: 1.1025 Accuracy: 0.8560 Time: 19.14605  | Val Loss: 0.7518 Accuracy: 0.8939\n",
      "Epoch: 44/100 Train Loss: 1.1012 Accuracy: 0.8554 Time: 18.78735  | Val Loss: 0.7588 Accuracy: 0.8906\n",
      "Epoch: 45/100 Train Loss: 1.0889 Accuracy: 0.8597 Time: 18.74859  | Val Loss: 0.7720 Accuracy: 0.8873\n",
      "Epoch: 46/100 Train Loss: 1.0863 Accuracy: 0.8613 Time: 19.00412  | Val Loss: 0.7410 Accuracy: 0.9021\n",
      "Epoch: 47/100 Train Loss: 1.0795 Accuracy: 0.8627 Time: 18.72991  | Val Loss: 0.7517 Accuracy: 0.8958\n",
      "Epoch: 48/100 Train Loss: 1.0704 Accuracy: 0.8646 Time: 18.57309  | Val Loss: 0.7455 Accuracy: 0.8979\n",
      "Epoch: 49/100 Train Loss: 1.0641 Accuracy: 0.8673 Time: 18.78999  | Val Loss: 0.7473 Accuracy: 0.8974\n",
      "Epoch: 50/100 Train Loss: 0.8204 Accuracy: 0.8604 Time: 17.36591  | Val Loss: 0.7442 Accuracy: 0.8984\n",
      "Epoch: 51/100 Train Loss: 0.8040 Accuracy: 0.8662 Time: 17.22856  | Val Loss: 0.7316 Accuracy: 0.9027\n",
      "Epoch: 52/100 Train Loss: 0.7882 Accuracy: 0.8749 Time: 17.00818  | Val Loss: 0.7406 Accuracy: 0.9000\n",
      "Epoch: 53/100 Train Loss: 0.7829 Accuracy: 0.8763 Time: 17.44288  | Val Loss: 0.7208 Accuracy: 0.9060\n",
      "Epoch: 54/100 Train Loss: 0.7792 Accuracy: 0.8788 Time: 17.30866  | Val Loss: 0.7151 Accuracy: 0.9100\n",
      "Epoch: 55/100 Train Loss: 0.7706 Accuracy: 0.8827 Time: 16.92018  | Val Loss: 0.7119 Accuracy: 0.9108\n",
      "Epoch: 56/100 Train Loss: 0.7607 Accuracy: 0.8861 Time: 17.01440  | Val Loss: 0.7117 Accuracy: 0.9111\n",
      "Epoch: 57/100 Train Loss: 0.7608 Accuracy: 0.8864 Time: 17.25801  | Val Loss: 0.7114 Accuracy: 0.9141\n",
      "Epoch: 58/100 Train Loss: 0.7522 Accuracy: 0.8902 Time: 17.34578  | Val Loss: 0.7010 Accuracy: 0.9154\n",
      "Epoch: 59/100 Train Loss: 0.7461 Accuracy: 0.8923 Time: 17.42741  | Val Loss: 0.7075 Accuracy: 0.9125\n",
      "Epoch: 60/100 Train Loss: 0.7482 Accuracy: 0.8933 Time: 17.57162  | Val Loss: 0.6977 Accuracy: 0.9166\n",
      "Epoch: 61/100 Train Loss: 0.7426 Accuracy: 0.8935 Time: 17.44152  | Val Loss: 0.7063 Accuracy: 0.9147\n",
      "Epoch: 62/100 Train Loss: 0.7378 Accuracy: 0.8972 Time: 17.17447  | Val Loss: 0.6979 Accuracy: 0.9161\n",
      "Epoch: 63/100 Train Loss: 0.7269 Accuracy: 0.9019 Time: 17.48883  | Val Loss: 0.7067 Accuracy: 0.9139\n",
      "Epoch: 64/100 Train Loss: 0.7279 Accuracy: 0.9012 Time: 17.28791  | Val Loss: 0.6972 Accuracy: 0.9155\n",
      "Epoch: 65/100 Train Loss: 0.7214 Accuracy: 0.9028 Time: 17.58077  | Val Loss: 0.6942 Accuracy: 0.9172\n",
      "Epoch: 66/100 Train Loss: 0.7184 Accuracy: 0.9053 Time: 17.72028  | Val Loss: 0.6867 Accuracy: 0.9228\n",
      "Epoch: 67/100 Train Loss: 0.7174 Accuracy: 0.9051 Time: 17.35684  | Val Loss: 0.6841 Accuracy: 0.9241\n",
      "Epoch: 68/100 Train Loss: 0.7132 Accuracy: 0.9073 Time: 17.65978  | Val Loss: 0.6870 Accuracy: 0.9231\n",
      "Epoch: 69/100 Train Loss: 0.7131 Accuracy: 0.9077 Time: 17.65487  | Val Loss: 0.6863 Accuracy: 0.9230\n",
      "Epoch: 70/100 Train Loss: 0.7072 Accuracy: 0.9099 Time: 17.83617  | Val Loss: 0.6792 Accuracy: 0.9251\n",
      "Epoch: 71/100 Train Loss: 0.7053 Accuracy: 0.9103 Time: 17.62511  | Val Loss: 0.6853 Accuracy: 0.9226\n",
      "Epoch: 72/100 Train Loss: 0.6982 Accuracy: 0.9148 Time: 17.50023  | Val Loss: 0.6860 Accuracy: 0.9221\n",
      "Epoch: 73/100 Train Loss: 0.6977 Accuracy: 0.9148 Time: 17.10662  | Val Loss: 0.6818 Accuracy: 0.9240\n",
      "Epoch: 74/100 Train Loss: 0.6973 Accuracy: 0.9139 Time: 17.59139  | Val Loss: 0.6826 Accuracy: 0.9223\n",
      "Epoch: 75/100 Train Loss: 0.6947 Accuracy: 0.9153 Time: 17.17783  | Val Loss: 0.6788 Accuracy: 0.9256\n",
      "Epoch: 76/100 Train Loss: 0.6874 Accuracy: 0.9195 Time: 16.92260  | Val Loss: 0.6757 Accuracy: 0.9256\n",
      "Epoch: 77/100 Train Loss: 0.6872 Accuracy: 0.9183 Time: 17.08025  | Val Loss: 0.6737 Accuracy: 0.9298\n",
      "Epoch: 78/100 Train Loss: 0.6834 Accuracy: 0.9196 Time: 17.08731  | Val Loss: 0.6810 Accuracy: 0.9261\n",
      "Epoch: 79/100 Train Loss: 0.6813 Accuracy: 0.9216 Time: 16.93729  | Val Loss: 0.6778 Accuracy: 0.9247\n",
      "Epoch: 80/100 Train Loss: 0.6807 Accuracy: 0.9218 Time: 17.68184  | Val Loss: 0.6776 Accuracy: 0.9254\n",
      "Epoch: 81/100 Train Loss: 0.6722 Accuracy: 0.9244 Time: 17.33746  | Val Loss: 0.6722 Accuracy: 0.9281\n",
      "Epoch: 82/100 Train Loss: 0.6748 Accuracy: 0.9238 Time: 16.91245  | Val Loss: 0.6753 Accuracy: 0.9251\n",
      "Epoch: 83/100 Train Loss: 0.6679 Accuracy: 0.9272 Time: 17.15582  | Val Loss: 0.6743 Accuracy: 0.9287\n",
      "Epoch: 84/100 Train Loss: 0.6672 Accuracy: 0.9277 Time: 17.13603  | Val Loss: 0.6715 Accuracy: 0.9304\n",
      "Epoch: 85/100 Train Loss: 0.6699 Accuracy: 0.9260 Time: 17.29221  | Val Loss: 0.6693 Accuracy: 0.9287\n",
      "Epoch: 86/100 Train Loss: 0.6676 Accuracy: 0.9281 Time: 17.31384  | Val Loss: 0.6731 Accuracy: 0.9287\n",
      "Epoch: 87/100 Train Loss: 0.6625 Accuracy: 0.9305 Time: 17.78837  | Val Loss: 0.6697 Accuracy: 0.9290\n",
      "Epoch: 88/100 Train Loss: 0.6615 Accuracy: 0.9292 Time: 17.13757  | Val Loss: 0.6668 Accuracy: 0.9309\n",
      "Epoch: 89/100 Train Loss: 0.6587 Accuracy: 0.9302 Time: 16.80385  | Val Loss: 0.6656 Accuracy: 0.9318\n",
      "Epoch: 90/100 Train Loss: 0.6574 Accuracy: 0.9325 Time: 17.49507  | Val Loss: 0.6718 Accuracy: 0.9325\n",
      "Epoch: 91/100 Train Loss: 0.6568 Accuracy: 0.9318 Time: 17.85246  | Val Loss: 0.6680 Accuracy: 0.9306\n",
      "Epoch: 92/100 Train Loss: 0.6563 Accuracy: 0.9322 Time: 17.32053  | Val Loss: 0.6661 Accuracy: 0.9325\n",
      "Epoch: 93/100 Train Loss: 0.6532 Accuracy: 0.9340 Time: 17.16480  | Val Loss: 0.6667 Accuracy: 0.9316\n",
      "Epoch: 94/100 Train Loss: 0.6516 Accuracy: 0.9345 Time: 17.12367  | Val Loss: 0.6656 Accuracy: 0.9309\n",
      "Epoch: 95/100 Train Loss: 0.6491 Accuracy: 0.9355 Time: 17.52183  | Val Loss: 0.6656 Accuracy: 0.9315\n",
      "Epoch: 96/100 Train Loss: 0.6493 Accuracy: 0.9349 Time: 17.85793  | Val Loss: 0.6668 Accuracy: 0.9330\n",
      "Epoch: 97/100 Train Loss: 0.6480 Accuracy: 0.9359 Time: 17.88172  | Val Loss: 0.6672 Accuracy: 0.9314\n",
      "Epoch: 98/100 Train Loss: 0.6483 Accuracy: 0.9357 Time: 18.37923  | Val Loss: 0.6678 Accuracy: 0.9321\n",
      "Epoch: 99/100 Train Loss: 0.6468 Accuracy: 0.9365 Time: 17.97911  | Val Loss: 0.6662 Accuracy: 0.9335\n",
      "Epoch: 100/100 Train Loss: 0.6441 Accuracy: 0.9374 Time: 17.20870  | Val Loss: 0.6671 Accuracy: 0.9311\n"
     ]
    }
   ],
   "source": [
    "from hdd.train.warmup_scheduler import GradualWarmupScheduler\n",
    "from hdd.models.cnn.condensenet import CondenseNet\n",
    "from hdd.models.nn_utils import count_trainable_parameter\n",
    "\n",
    "\n",
    "net = CondenseNet(\n",
    "    num_classes=10,\n",
    "    group_1x1=4,\n",
    "    group_3x3=4,\n",
    "    bottleneck=4,\n",
    "    condense_factor=4,\n",
    "    dropout_rate=0,\n",
    "    stages=[14, 14, 14],\n",
    "    growth=[8, 16, 32],\n",
    "    data=\"cifar10\",\n",
    ").to(DEVICE)\n",
    "print(f\"#Parameter: {count_trainable_parameter(net)}\")\n",
    "criteria = nn.CrossEntropyLoss(label_smoothing=0.1)\n",
    "optimizer = torch.optim.Adam(\n",
    "    net.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=1e-5\n",
    ")\n",
    "max_epochs = 100\n",
    "base_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
    "    optimizer, max_epochs, eta_min=1e-5\n",
    ")\n",
    "scheduler = GradualWarmupScheduler(\n",
    "    optimizer,\n",
    "    multiplier=1.0,\n",
    "    total_epoch=10,\n",
    "    after_scheduler=base_scheduler,\n",
    ")\n",
    "lasso_lambda = 1e-5\n",
    "random_setting = train_condensenet_model(\n",
    "    net,\n",
    "    criteria,\n",
    "    max_epochs,\n",
    "    lasso_lambda,\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    DEVICE,\n",
    "    optimizer,\n",
    "    scheduler,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "410bd41d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1           [-1, 16, 32, 32]             432\n",
      "       BatchNorm2d-2           [-1, 16, 32, 32]              32\n",
      "              ReLU-3           [-1, 16, 32, 32]               0\n",
      "  LearnedGroupConv-4           [-1, 32, 32, 32]               0\n",
      "       BatchNorm2d-5           [-1, 32, 32, 32]              64\n",
      "              ReLU-6           [-1, 32, 32, 32]               0\n",
      "            Conv2d-7            [-1, 8, 32, 32]             576\n",
      "       _DenseLayer-8           [-1, 24, 32, 32]               0\n",
      "       BatchNorm2d-9           [-1, 24, 32, 32]              48\n",
      "             ReLU-10           [-1, 24, 32, 32]               0\n",
      " LearnedGroupConv-11           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-12           [-1, 32, 32, 32]              64\n",
      "             ReLU-13           [-1, 32, 32, 32]               0\n",
      "           Conv2d-14            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-15           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-16           [-1, 32, 32, 32]              64\n",
      "             ReLU-17           [-1, 32, 32, 32]               0\n",
      " LearnedGroupConv-18           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-19           [-1, 32, 32, 32]              64\n",
      "             ReLU-20           [-1, 32, 32, 32]               0\n",
      "           Conv2d-21            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-22           [-1, 40, 32, 32]               0\n",
      "      BatchNorm2d-23           [-1, 40, 32, 32]              80\n",
      "             ReLU-24           [-1, 40, 32, 32]               0\n",
      " LearnedGroupConv-25           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-26           [-1, 32, 32, 32]              64\n",
      "             ReLU-27           [-1, 32, 32, 32]               0\n",
      "           Conv2d-28            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-29           [-1, 48, 32, 32]               0\n",
      "      BatchNorm2d-30           [-1, 48, 32, 32]              96\n",
      "             ReLU-31           [-1, 48, 32, 32]               0\n",
      " LearnedGroupConv-32           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-33           [-1, 32, 32, 32]              64\n",
      "             ReLU-34           [-1, 32, 32, 32]               0\n",
      "           Conv2d-35            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-36           [-1, 56, 32, 32]               0\n",
      "      BatchNorm2d-37           [-1, 56, 32, 32]             112\n",
      "             ReLU-38           [-1, 56, 32, 32]               0\n",
      " LearnedGroupConv-39           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-40           [-1, 32, 32, 32]              64\n",
      "             ReLU-41           [-1, 32, 32, 32]               0\n",
      "           Conv2d-42            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-43           [-1, 64, 32, 32]               0\n",
      "      BatchNorm2d-44           [-1, 64, 32, 32]             128\n",
      "             ReLU-45           [-1, 64, 32, 32]               0\n",
      " LearnedGroupConv-46           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-47           [-1, 32, 32, 32]              64\n",
      "             ReLU-48           [-1, 32, 32, 32]               0\n",
      "           Conv2d-49            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-50           [-1, 72, 32, 32]               0\n",
      "      BatchNorm2d-51           [-1, 72, 32, 32]             144\n",
      "             ReLU-52           [-1, 72, 32, 32]               0\n",
      " LearnedGroupConv-53           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-54           [-1, 32, 32, 32]              64\n",
      "             ReLU-55           [-1, 32, 32, 32]               0\n",
      "           Conv2d-56            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-57           [-1, 80, 32, 32]               0\n",
      "      BatchNorm2d-58           [-1, 80, 32, 32]             160\n",
      "             ReLU-59           [-1, 80, 32, 32]               0\n",
      " LearnedGroupConv-60           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-61           [-1, 32, 32, 32]              64\n",
      "             ReLU-62           [-1, 32, 32, 32]               0\n",
      "           Conv2d-63            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-64           [-1, 88, 32, 32]               0\n",
      "      BatchNorm2d-65           [-1, 88, 32, 32]             176\n",
      "             ReLU-66           [-1, 88, 32, 32]               0\n",
      " LearnedGroupConv-67           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-68           [-1, 32, 32, 32]              64\n",
      "             ReLU-69           [-1, 32, 32, 32]               0\n",
      "           Conv2d-70            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-71           [-1, 96, 32, 32]               0\n",
      "      BatchNorm2d-72           [-1, 96, 32, 32]             192\n",
      "             ReLU-73           [-1, 96, 32, 32]               0\n",
      " LearnedGroupConv-74           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-75           [-1, 32, 32, 32]              64\n",
      "             ReLU-76           [-1, 32, 32, 32]               0\n",
      "           Conv2d-77            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-78          [-1, 104, 32, 32]               0\n",
      "      BatchNorm2d-79          [-1, 104, 32, 32]             208\n",
      "             ReLU-80          [-1, 104, 32, 32]               0\n",
      " LearnedGroupConv-81           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-82           [-1, 32, 32, 32]              64\n",
      "             ReLU-83           [-1, 32, 32, 32]               0\n",
      "           Conv2d-84            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-85          [-1, 112, 32, 32]               0\n",
      "      BatchNorm2d-86          [-1, 112, 32, 32]             224\n",
      "             ReLU-87          [-1, 112, 32, 32]               0\n",
      " LearnedGroupConv-88           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-89           [-1, 32, 32, 32]              64\n",
      "             ReLU-90           [-1, 32, 32, 32]               0\n",
      "           Conv2d-91            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-92          [-1, 120, 32, 32]               0\n",
      "      BatchNorm2d-93          [-1, 120, 32, 32]             240\n",
      "             ReLU-94          [-1, 120, 32, 32]               0\n",
      " LearnedGroupConv-95           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-96           [-1, 32, 32, 32]              64\n",
      "             ReLU-97           [-1, 32, 32, 32]               0\n",
      "           Conv2d-98            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-99          [-1, 128, 32, 32]               0\n",
      "       AvgPool2d-100          [-1, 128, 16, 16]               0\n",
      "     _Transition-101          [-1, 128, 16, 16]               0\n",
      "     BatchNorm2d-102          [-1, 128, 16, 16]             256\n",
      "            ReLU-103          [-1, 128, 16, 16]               0\n",
      "LearnedGroupConv-104           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-105           [-1, 64, 16, 16]             128\n",
      "            ReLU-106           [-1, 64, 16, 16]               0\n",
      "          Conv2d-107           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-108          [-1, 144, 16, 16]               0\n",
      "     BatchNorm2d-109          [-1, 144, 16, 16]             288\n",
      "            ReLU-110          [-1, 144, 16, 16]               0\n",
      "LearnedGroupConv-111           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-112           [-1, 64, 16, 16]             128\n",
      "            ReLU-113           [-1, 64, 16, 16]               0\n",
      "          Conv2d-114           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-115          [-1, 160, 16, 16]               0\n",
      "     BatchNorm2d-116          [-1, 160, 16, 16]             320\n",
      "            ReLU-117          [-1, 160, 16, 16]               0\n",
      "LearnedGroupConv-118           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-119           [-1, 64, 16, 16]             128\n",
      "            ReLU-120           [-1, 64, 16, 16]               0\n",
      "          Conv2d-121           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-122          [-1, 176, 16, 16]               0\n",
      "     BatchNorm2d-123          [-1, 176, 16, 16]             352\n",
      "            ReLU-124          [-1, 176, 16, 16]               0\n",
      "LearnedGroupConv-125           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-126           [-1, 64, 16, 16]             128\n",
      "            ReLU-127           [-1, 64, 16, 16]               0\n",
      "          Conv2d-128           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-129          [-1, 192, 16, 16]               0\n",
      "     BatchNorm2d-130          [-1, 192, 16, 16]             384\n",
      "            ReLU-131          [-1, 192, 16, 16]               0\n",
      "LearnedGroupConv-132           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-133           [-1, 64, 16, 16]             128\n",
      "            ReLU-134           [-1, 64, 16, 16]               0\n",
      "          Conv2d-135           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-136          [-1, 208, 16, 16]               0\n",
      "     BatchNorm2d-137          [-1, 208, 16, 16]             416\n",
      "            ReLU-138          [-1, 208, 16, 16]               0\n",
      "LearnedGroupConv-139           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-140           [-1, 64, 16, 16]             128\n",
      "            ReLU-141           [-1, 64, 16, 16]               0\n",
      "          Conv2d-142           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-143          [-1, 224, 16, 16]               0\n",
      "     BatchNorm2d-144          [-1, 224, 16, 16]             448\n",
      "            ReLU-145          [-1, 224, 16, 16]               0\n",
      "LearnedGroupConv-146           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-147           [-1, 64, 16, 16]             128\n",
      "            ReLU-148           [-1, 64, 16, 16]               0\n",
      "          Conv2d-149           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-150          [-1, 240, 16, 16]               0\n",
      "     BatchNorm2d-151          [-1, 240, 16, 16]             480\n",
      "            ReLU-152          [-1, 240, 16, 16]               0\n",
      "LearnedGroupConv-153           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-154           [-1, 64, 16, 16]             128\n",
      "            ReLU-155           [-1, 64, 16, 16]               0\n",
      "          Conv2d-156           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-157          [-1, 256, 16, 16]               0\n",
      "     BatchNorm2d-158          [-1, 256, 16, 16]             512\n",
      "            ReLU-159          [-1, 256, 16, 16]               0\n",
      "LearnedGroupConv-160           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-161           [-1, 64, 16, 16]             128\n",
      "            ReLU-162           [-1, 64, 16, 16]               0\n",
      "          Conv2d-163           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-164          [-1, 272, 16, 16]               0\n",
      "     BatchNorm2d-165          [-1, 272, 16, 16]             544\n",
      "            ReLU-166          [-1, 272, 16, 16]               0\n",
      "LearnedGroupConv-167           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-168           [-1, 64, 16, 16]             128\n",
      "            ReLU-169           [-1, 64, 16, 16]               0\n",
      "          Conv2d-170           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-171          [-1, 288, 16, 16]               0\n",
      "     BatchNorm2d-172          [-1, 288, 16, 16]             576\n",
      "            ReLU-173          [-1, 288, 16, 16]               0\n",
      "LearnedGroupConv-174           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-175           [-1, 64, 16, 16]             128\n",
      "            ReLU-176           [-1, 64, 16, 16]               0\n",
      "          Conv2d-177           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-178          [-1, 304, 16, 16]               0\n",
      "     BatchNorm2d-179          [-1, 304, 16, 16]             608\n",
      "            ReLU-180          [-1, 304, 16, 16]               0\n",
      "LearnedGroupConv-181           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-182           [-1, 64, 16, 16]             128\n",
      "            ReLU-183           [-1, 64, 16, 16]               0\n",
      "          Conv2d-184           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-185          [-1, 320, 16, 16]               0\n",
      "     BatchNorm2d-186          [-1, 320, 16, 16]             640\n",
      "            ReLU-187          [-1, 320, 16, 16]               0\n",
      "LearnedGroupConv-188           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-189           [-1, 64, 16, 16]             128\n",
      "            ReLU-190           [-1, 64, 16, 16]               0\n",
      "          Conv2d-191           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-192          [-1, 336, 16, 16]               0\n",
      "     BatchNorm2d-193          [-1, 336, 16, 16]             672\n",
      "            ReLU-194          [-1, 336, 16, 16]               0\n",
      "LearnedGroupConv-195           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-196           [-1, 64, 16, 16]             128\n",
      "            ReLU-197           [-1, 64, 16, 16]               0\n",
      "          Conv2d-198           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-199          [-1, 352, 16, 16]               0\n",
      "       AvgPool2d-200            [-1, 352, 8, 8]               0\n",
      "     _Transition-201            [-1, 352, 8, 8]               0\n",
      "     BatchNorm2d-202            [-1, 352, 8, 8]             704\n",
      "            ReLU-203            [-1, 352, 8, 8]               0\n",
      "LearnedGroupConv-204            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-205            [-1, 128, 8, 8]             256\n",
      "            ReLU-206            [-1, 128, 8, 8]               0\n",
      "          Conv2d-207             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-208            [-1, 384, 8, 8]               0\n",
      "     BatchNorm2d-209            [-1, 384, 8, 8]             768\n",
      "            ReLU-210            [-1, 384, 8, 8]               0\n",
      "LearnedGroupConv-211            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-212            [-1, 128, 8, 8]             256\n",
      "            ReLU-213            [-1, 128, 8, 8]               0\n",
      "          Conv2d-214             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-215            [-1, 416, 8, 8]               0\n",
      "     BatchNorm2d-216            [-1, 416, 8, 8]             832\n",
      "            ReLU-217            [-1, 416, 8, 8]               0\n",
      "LearnedGroupConv-218            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-219            [-1, 128, 8, 8]             256\n",
      "            ReLU-220            [-1, 128, 8, 8]               0\n",
      "          Conv2d-221             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-222            [-1, 448, 8, 8]               0\n",
      "     BatchNorm2d-223            [-1, 448, 8, 8]             896\n",
      "            ReLU-224            [-1, 448, 8, 8]               0\n",
      "LearnedGroupConv-225            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-226            [-1, 128, 8, 8]             256\n",
      "            ReLU-227            [-1, 128, 8, 8]               0\n",
      "          Conv2d-228             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-229            [-1, 480, 8, 8]               0\n",
      "     BatchNorm2d-230            [-1, 480, 8, 8]             960\n",
      "            ReLU-231            [-1, 480, 8, 8]               0\n",
      "LearnedGroupConv-232            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-233            [-1, 128, 8, 8]             256\n",
      "            ReLU-234            [-1, 128, 8, 8]               0\n",
      "          Conv2d-235             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-236            [-1, 512, 8, 8]               0\n",
      "     BatchNorm2d-237            [-1, 512, 8, 8]           1,024\n",
      "            ReLU-238            [-1, 512, 8, 8]               0\n",
      "LearnedGroupConv-239            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-240            [-1, 128, 8, 8]             256\n",
      "            ReLU-241            [-1, 128, 8, 8]               0\n",
      "          Conv2d-242             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-243            [-1, 544, 8, 8]               0\n",
      "     BatchNorm2d-244            [-1, 544, 8, 8]           1,088\n",
      "            ReLU-245            [-1, 544, 8, 8]               0\n",
      "LearnedGroupConv-246            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-247            [-1, 128, 8, 8]             256\n",
      "            ReLU-248            [-1, 128, 8, 8]               0\n",
      "          Conv2d-249             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-250            [-1, 576, 8, 8]               0\n",
      "     BatchNorm2d-251            [-1, 576, 8, 8]           1,152\n",
      "            ReLU-252            [-1, 576, 8, 8]               0\n",
      "LearnedGroupConv-253            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-254            [-1, 128, 8, 8]             256\n",
      "            ReLU-255            [-1, 128, 8, 8]               0\n",
      "          Conv2d-256             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-257            [-1, 608, 8, 8]               0\n",
      "     BatchNorm2d-258            [-1, 608, 8, 8]           1,216\n",
      "            ReLU-259            [-1, 608, 8, 8]               0\n",
      "LearnedGroupConv-260            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-261            [-1, 128, 8, 8]             256\n",
      "            ReLU-262            [-1, 128, 8, 8]               0\n",
      "          Conv2d-263             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-264            [-1, 640, 8, 8]               0\n",
      "     BatchNorm2d-265            [-1, 640, 8, 8]           1,280\n",
      "            ReLU-266            [-1, 640, 8, 8]               0\n",
      "LearnedGroupConv-267            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-268            [-1, 128, 8, 8]             256\n",
      "            ReLU-269            [-1, 128, 8, 8]               0\n",
      "          Conv2d-270             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-271            [-1, 672, 8, 8]               0\n",
      "     BatchNorm2d-272            [-1, 672, 8, 8]           1,344\n",
      "            ReLU-273            [-1, 672, 8, 8]               0\n",
      "LearnedGroupConv-274            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-275            [-1, 128, 8, 8]             256\n",
      "            ReLU-276            [-1, 128, 8, 8]               0\n",
      "          Conv2d-277             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-278            [-1, 704, 8, 8]               0\n",
      "     BatchNorm2d-279            [-1, 704, 8, 8]           1,408\n",
      "            ReLU-280            [-1, 704, 8, 8]               0\n",
      "LearnedGroupConv-281            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-282            [-1, 128, 8, 8]             256\n",
      "            ReLU-283            [-1, 128, 8, 8]               0\n",
      "          Conv2d-284             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-285            [-1, 736, 8, 8]               0\n",
      "     BatchNorm2d-286            [-1, 736, 8, 8]           1,472\n",
      "            ReLU-287            [-1, 736, 8, 8]               0\n",
      "LearnedGroupConv-288            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-289            [-1, 128, 8, 8]             256\n",
      "            ReLU-290            [-1, 128, 8, 8]               0\n",
      "          Conv2d-291             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-292            [-1, 768, 8, 8]               0\n",
      "     BatchNorm2d-293            [-1, 768, 8, 8]           1,536\n",
      "            ReLU-294            [-1, 768, 8, 8]               0\n",
      "LearnedGroupConv-295            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-296            [-1, 128, 8, 8]             256\n",
      "            ReLU-297            [-1, 128, 8, 8]               0\n",
      "          Conv2d-298             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-299            [-1, 800, 8, 8]               0\n",
      "     BatchNorm2d-300            [-1, 800, 8, 8]           1,600\n",
      "            ReLU-301            [-1, 800, 8, 8]               0\n",
      "       AvgPool2d-302            [-1, 800, 1, 1]               0\n",
      "          Linear-303                   [-1, 10]           8,010\n",
      "================================================================\n",
      "Total params: 209,738\n",
      "Trainable params: 209,738\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.01\n",
      "Forward/backward pass size (MB): 76.02\n",
      "Params size (MB): 0.80\n",
      "Estimated Total Size (MB): 76.83\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "from torchsummary import summary\n",
    "\n",
    "summary(net, (3, 32, 32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d562e3bc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1           [-1, 16, 32, 32]             432\n",
      "       BatchNorm2d-2           [-1, 16, 32, 32]              32\n",
      "              ReLU-3           [-1, 16, 32, 32]               0\n",
      "            Conv2d-4           [-1, 32, 32, 32]             128\n",
      "    CondensingConv-5           [-1, 32, 32, 32]               0\n",
      "       BatchNorm2d-6           [-1, 32, 32, 32]              64\n",
      "              ReLU-7           [-1, 32, 32, 32]               0\n",
      "            Conv2d-8            [-1, 8, 32, 32]             576\n",
      "       _DenseLayer-9           [-1, 24, 32, 32]               0\n",
      "      BatchNorm2d-10           [-1, 24, 32, 32]              48\n",
      "             ReLU-11           [-1, 24, 32, 32]               0\n",
      "           Conv2d-12           [-1, 32, 32, 32]             192\n",
      "   CondensingConv-13           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-14           [-1, 32, 32, 32]              64\n",
      "             ReLU-15           [-1, 32, 32, 32]               0\n",
      "           Conv2d-16            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-17           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-18           [-1, 32, 32, 32]              64\n",
      "             ReLU-19           [-1, 32, 32, 32]               0\n",
      "           Conv2d-20           [-1, 32, 32, 32]             256\n",
      "   CondensingConv-21           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-22           [-1, 32, 32, 32]              64\n",
      "             ReLU-23           [-1, 32, 32, 32]               0\n",
      "           Conv2d-24            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-25           [-1, 40, 32, 32]               0\n",
      "      BatchNorm2d-26           [-1, 40, 32, 32]              80\n",
      "             ReLU-27           [-1, 40, 32, 32]               0\n",
      "           Conv2d-28           [-1, 32, 32, 32]             320\n",
      "   CondensingConv-29           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-30           [-1, 32, 32, 32]              64\n",
      "             ReLU-31           [-1, 32, 32, 32]               0\n",
      "           Conv2d-32            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-33           [-1, 48, 32, 32]               0\n",
      "      BatchNorm2d-34           [-1, 48, 32, 32]              96\n",
      "             ReLU-35           [-1, 48, 32, 32]               0\n",
      "           Conv2d-36           [-1, 32, 32, 32]             384\n",
      "   CondensingConv-37           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-38           [-1, 32, 32, 32]              64\n",
      "             ReLU-39           [-1, 32, 32, 32]               0\n",
      "           Conv2d-40            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-41           [-1, 56, 32, 32]               0\n",
      "      BatchNorm2d-42           [-1, 56, 32, 32]             112\n",
      "             ReLU-43           [-1, 56, 32, 32]               0\n",
      "           Conv2d-44           [-1, 32, 32, 32]             448\n",
      "   CondensingConv-45           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-46           [-1, 32, 32, 32]              64\n",
      "             ReLU-47           [-1, 32, 32, 32]               0\n",
      "           Conv2d-48            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-49           [-1, 64, 32, 32]               0\n",
      "      BatchNorm2d-50           [-1, 64, 32, 32]             128\n",
      "             ReLU-51           [-1, 64, 32, 32]               0\n",
      "           Conv2d-52           [-1, 32, 32, 32]             512\n",
      "   CondensingConv-53           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-54           [-1, 32, 32, 32]              64\n",
      "             ReLU-55           [-1, 32, 32, 32]               0\n",
      "           Conv2d-56            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-57           [-1, 72, 32, 32]               0\n",
      "      BatchNorm2d-58           [-1, 72, 32, 32]             144\n",
      "             ReLU-59           [-1, 72, 32, 32]               0\n",
      "           Conv2d-60           [-1, 32, 32, 32]             576\n",
      "   CondensingConv-61           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-62           [-1, 32, 32, 32]              64\n",
      "             ReLU-63           [-1, 32, 32, 32]               0\n",
      "           Conv2d-64            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-65           [-1, 80, 32, 32]               0\n",
      "      BatchNorm2d-66           [-1, 80, 32, 32]             160\n",
      "             ReLU-67           [-1, 80, 32, 32]               0\n",
      "           Conv2d-68           [-1, 32, 32, 32]             640\n",
      "   CondensingConv-69           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-70           [-1, 32, 32, 32]              64\n",
      "             ReLU-71           [-1, 32, 32, 32]               0\n",
      "           Conv2d-72            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-73           [-1, 88, 32, 32]               0\n",
      "      BatchNorm2d-74           [-1, 88, 32, 32]             176\n",
      "             ReLU-75           [-1, 88, 32, 32]               0\n",
      "           Conv2d-76           [-1, 32, 32, 32]             704\n",
      "   CondensingConv-77           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-78           [-1, 32, 32, 32]              64\n",
      "             ReLU-79           [-1, 32, 32, 32]               0\n",
      "           Conv2d-80            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-81           [-1, 96, 32, 32]               0\n",
      "      BatchNorm2d-82           [-1, 96, 32, 32]             192\n",
      "             ReLU-83           [-1, 96, 32, 32]               0\n",
      "           Conv2d-84           [-1, 32, 32, 32]             768\n",
      "   CondensingConv-85           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-86           [-1, 32, 32, 32]              64\n",
      "             ReLU-87           [-1, 32, 32, 32]               0\n",
      "           Conv2d-88            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-89          [-1, 104, 32, 32]               0\n",
      "      BatchNorm2d-90          [-1, 104, 32, 32]             208\n",
      "             ReLU-91          [-1, 104, 32, 32]               0\n",
      "           Conv2d-92           [-1, 32, 32, 32]             832\n",
      "   CondensingConv-93           [-1, 32, 32, 32]               0\n",
      "      BatchNorm2d-94           [-1, 32, 32, 32]              64\n",
      "             ReLU-95           [-1, 32, 32, 32]               0\n",
      "           Conv2d-96            [-1, 8, 32, 32]             576\n",
      "      _DenseLayer-97          [-1, 112, 32, 32]               0\n",
      "      BatchNorm2d-98          [-1, 112, 32, 32]             224\n",
      "             ReLU-99          [-1, 112, 32, 32]               0\n",
      "          Conv2d-100           [-1, 32, 32, 32]             896\n",
      "  CondensingConv-101           [-1, 32, 32, 32]               0\n",
      "     BatchNorm2d-102           [-1, 32, 32, 32]              64\n",
      "            ReLU-103           [-1, 32, 32, 32]               0\n",
      "          Conv2d-104            [-1, 8, 32, 32]             576\n",
      "     _DenseLayer-105          [-1, 120, 32, 32]               0\n",
      "     BatchNorm2d-106          [-1, 120, 32, 32]             240\n",
      "            ReLU-107          [-1, 120, 32, 32]               0\n",
      "          Conv2d-108           [-1, 32, 32, 32]             960\n",
      "  CondensingConv-109           [-1, 32, 32, 32]               0\n",
      "     BatchNorm2d-110           [-1, 32, 32, 32]              64\n",
      "            ReLU-111           [-1, 32, 32, 32]               0\n",
      "          Conv2d-112            [-1, 8, 32, 32]             576\n",
      "     _DenseLayer-113          [-1, 128, 32, 32]               0\n",
      "       AvgPool2d-114          [-1, 128, 16, 16]               0\n",
      "     _Transition-115          [-1, 128, 16, 16]               0\n",
      "     BatchNorm2d-116          [-1, 128, 16, 16]             256\n",
      "            ReLU-117          [-1, 128, 16, 16]               0\n",
      "          Conv2d-118           [-1, 64, 16, 16]           2,048\n",
      "  CondensingConv-119           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-120           [-1, 64, 16, 16]             128\n",
      "            ReLU-121           [-1, 64, 16, 16]               0\n",
      "          Conv2d-122           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-123          [-1, 144, 16, 16]               0\n",
      "     BatchNorm2d-124          [-1, 144, 16, 16]             288\n",
      "            ReLU-125          [-1, 144, 16, 16]               0\n",
      "          Conv2d-126           [-1, 64, 16, 16]           2,304\n",
      "  CondensingConv-127           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-128           [-1, 64, 16, 16]             128\n",
      "            ReLU-129           [-1, 64, 16, 16]               0\n",
      "          Conv2d-130           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-131          [-1, 160, 16, 16]               0\n",
      "     BatchNorm2d-132          [-1, 160, 16, 16]             320\n",
      "            ReLU-133          [-1, 160, 16, 16]               0\n",
      "          Conv2d-134           [-1, 64, 16, 16]           2,560\n",
      "  CondensingConv-135           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-136           [-1, 64, 16, 16]             128\n",
      "            ReLU-137           [-1, 64, 16, 16]               0\n",
      "          Conv2d-138           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-139          [-1, 176, 16, 16]               0\n",
      "     BatchNorm2d-140          [-1, 176, 16, 16]             352\n",
      "            ReLU-141          [-1, 176, 16, 16]               0\n",
      "          Conv2d-142           [-1, 64, 16, 16]           2,816\n",
      "  CondensingConv-143           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-144           [-1, 64, 16, 16]             128\n",
      "            ReLU-145           [-1, 64, 16, 16]               0\n",
      "          Conv2d-146           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-147          [-1, 192, 16, 16]               0\n",
      "     BatchNorm2d-148          [-1, 192, 16, 16]             384\n",
      "            ReLU-149          [-1, 192, 16, 16]               0\n",
      "          Conv2d-150           [-1, 64, 16, 16]           3,072\n",
      "  CondensingConv-151           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-152           [-1, 64, 16, 16]             128\n",
      "            ReLU-153           [-1, 64, 16, 16]               0\n",
      "          Conv2d-154           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-155          [-1, 208, 16, 16]               0\n",
      "     BatchNorm2d-156          [-1, 208, 16, 16]             416\n",
      "            ReLU-157          [-1, 208, 16, 16]               0\n",
      "          Conv2d-158           [-1, 64, 16, 16]           3,328\n",
      "  CondensingConv-159           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-160           [-1, 64, 16, 16]             128\n",
      "            ReLU-161           [-1, 64, 16, 16]               0\n",
      "          Conv2d-162           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-163          [-1, 224, 16, 16]               0\n",
      "     BatchNorm2d-164          [-1, 224, 16, 16]             448\n",
      "            ReLU-165          [-1, 224, 16, 16]               0\n",
      "          Conv2d-166           [-1, 64, 16, 16]           3,584\n",
      "  CondensingConv-167           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-168           [-1, 64, 16, 16]             128\n",
      "            ReLU-169           [-1, 64, 16, 16]               0\n",
      "          Conv2d-170           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-171          [-1, 240, 16, 16]               0\n",
      "     BatchNorm2d-172          [-1, 240, 16, 16]             480\n",
      "            ReLU-173          [-1, 240, 16, 16]               0\n",
      "          Conv2d-174           [-1, 64, 16, 16]           3,840\n",
      "  CondensingConv-175           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-176           [-1, 64, 16, 16]             128\n",
      "            ReLU-177           [-1, 64, 16, 16]               0\n",
      "          Conv2d-178           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-179          [-1, 256, 16, 16]               0\n",
      "     BatchNorm2d-180          [-1, 256, 16, 16]             512\n",
      "            ReLU-181          [-1, 256, 16, 16]               0\n",
      "          Conv2d-182           [-1, 64, 16, 16]           4,096\n",
      "  CondensingConv-183           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-184           [-1, 64, 16, 16]             128\n",
      "            ReLU-185           [-1, 64, 16, 16]               0\n",
      "          Conv2d-186           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-187          [-1, 272, 16, 16]               0\n",
      "     BatchNorm2d-188          [-1, 272, 16, 16]             544\n",
      "            ReLU-189          [-1, 272, 16, 16]               0\n",
      "          Conv2d-190           [-1, 64, 16, 16]           4,352\n",
      "  CondensingConv-191           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-192           [-1, 64, 16, 16]             128\n",
      "            ReLU-193           [-1, 64, 16, 16]               0\n",
      "          Conv2d-194           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-195          [-1, 288, 16, 16]               0\n",
      "     BatchNorm2d-196          [-1, 288, 16, 16]             576\n",
      "            ReLU-197          [-1, 288, 16, 16]               0\n",
      "          Conv2d-198           [-1, 64, 16, 16]           4,608\n",
      "  CondensingConv-199           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-200           [-1, 64, 16, 16]             128\n",
      "            ReLU-201           [-1, 64, 16, 16]               0\n",
      "          Conv2d-202           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-203          [-1, 304, 16, 16]               0\n",
      "     BatchNorm2d-204          [-1, 304, 16, 16]             608\n",
      "            ReLU-205          [-1, 304, 16, 16]               0\n",
      "          Conv2d-206           [-1, 64, 16, 16]           4,864\n",
      "  CondensingConv-207           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-208           [-1, 64, 16, 16]             128\n",
      "            ReLU-209           [-1, 64, 16, 16]               0\n",
      "          Conv2d-210           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-211          [-1, 320, 16, 16]               0\n",
      "     BatchNorm2d-212          [-1, 320, 16, 16]             640\n",
      "            ReLU-213          [-1, 320, 16, 16]               0\n",
      "          Conv2d-214           [-1, 64, 16, 16]           5,120\n",
      "  CondensingConv-215           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-216           [-1, 64, 16, 16]             128\n",
      "            ReLU-217           [-1, 64, 16, 16]               0\n",
      "          Conv2d-218           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-219          [-1, 336, 16, 16]               0\n",
      "     BatchNorm2d-220          [-1, 336, 16, 16]             672\n",
      "            ReLU-221          [-1, 336, 16, 16]               0\n",
      "          Conv2d-222           [-1, 64, 16, 16]           5,376\n",
      "  CondensingConv-223           [-1, 64, 16, 16]               0\n",
      "     BatchNorm2d-224           [-1, 64, 16, 16]             128\n",
      "            ReLU-225           [-1, 64, 16, 16]               0\n",
      "          Conv2d-226           [-1, 16, 16, 16]           2,304\n",
      "     _DenseLayer-227          [-1, 352, 16, 16]               0\n",
      "       AvgPool2d-228            [-1, 352, 8, 8]               0\n",
      "     _Transition-229            [-1, 352, 8, 8]               0\n",
      "     BatchNorm2d-230            [-1, 352, 8, 8]             704\n",
      "            ReLU-231            [-1, 352, 8, 8]               0\n",
      "          Conv2d-232            [-1, 128, 8, 8]          11,264\n",
      "  CondensingConv-233            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-234            [-1, 128, 8, 8]             256\n",
      "            ReLU-235            [-1, 128, 8, 8]               0\n",
      "          Conv2d-236             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-237            [-1, 384, 8, 8]               0\n",
      "     BatchNorm2d-238            [-1, 384, 8, 8]             768\n",
      "            ReLU-239            [-1, 384, 8, 8]               0\n",
      "          Conv2d-240            [-1, 128, 8, 8]          12,288\n",
      "  CondensingConv-241            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-242            [-1, 128, 8, 8]             256\n",
      "            ReLU-243            [-1, 128, 8, 8]               0\n",
      "          Conv2d-244             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-245            [-1, 416, 8, 8]               0\n",
      "     BatchNorm2d-246            [-1, 416, 8, 8]             832\n",
      "            ReLU-247            [-1, 416, 8, 8]               0\n",
      "          Conv2d-248            [-1, 128, 8, 8]          13,312\n",
      "  CondensingConv-249            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-250            [-1, 128, 8, 8]             256\n",
      "            ReLU-251            [-1, 128, 8, 8]               0\n",
      "          Conv2d-252             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-253            [-1, 448, 8, 8]               0\n",
      "     BatchNorm2d-254            [-1, 448, 8, 8]             896\n",
      "            ReLU-255            [-1, 448, 8, 8]               0\n",
      "          Conv2d-256            [-1, 128, 8, 8]          14,336\n",
      "  CondensingConv-257            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-258            [-1, 128, 8, 8]             256\n",
      "            ReLU-259            [-1, 128, 8, 8]               0\n",
      "          Conv2d-260             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-261            [-1, 480, 8, 8]               0\n",
      "     BatchNorm2d-262            [-1, 480, 8, 8]             960\n",
      "            ReLU-263            [-1, 480, 8, 8]               0\n",
      "          Conv2d-264            [-1, 128, 8, 8]          15,360\n",
      "  CondensingConv-265            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-266            [-1, 128, 8, 8]             256\n",
      "            ReLU-267            [-1, 128, 8, 8]               0\n",
      "          Conv2d-268             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-269            [-1, 512, 8, 8]               0\n",
      "     BatchNorm2d-270            [-1, 512, 8, 8]           1,024\n",
      "            ReLU-271            [-1, 512, 8, 8]               0\n",
      "          Conv2d-272            [-1, 128, 8, 8]          16,384\n",
      "  CondensingConv-273            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-274            [-1, 128, 8, 8]             256\n",
      "            ReLU-275            [-1, 128, 8, 8]               0\n",
      "          Conv2d-276             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-277            [-1, 544, 8, 8]               0\n",
      "     BatchNorm2d-278            [-1, 544, 8, 8]           1,088\n",
      "            ReLU-279            [-1, 544, 8, 8]               0\n",
      "          Conv2d-280            [-1, 128, 8, 8]          17,408\n",
      "  CondensingConv-281            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-282            [-1, 128, 8, 8]             256\n",
      "            ReLU-283            [-1, 128, 8, 8]               0\n",
      "          Conv2d-284             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-285            [-1, 576, 8, 8]               0\n",
      "     BatchNorm2d-286            [-1, 576, 8, 8]           1,152\n",
      "            ReLU-287            [-1, 576, 8, 8]               0\n",
      "          Conv2d-288            [-1, 128, 8, 8]          18,432\n",
      "  CondensingConv-289            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-290            [-1, 128, 8, 8]             256\n",
      "            ReLU-291            [-1, 128, 8, 8]               0\n",
      "          Conv2d-292             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-293            [-1, 608, 8, 8]               0\n",
      "     BatchNorm2d-294            [-1, 608, 8, 8]           1,216\n",
      "            ReLU-295            [-1, 608, 8, 8]               0\n",
      "          Conv2d-296            [-1, 128, 8, 8]          19,456\n",
      "  CondensingConv-297            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-298            [-1, 128, 8, 8]             256\n",
      "            ReLU-299            [-1, 128, 8, 8]               0\n",
      "          Conv2d-300             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-301            [-1, 640, 8, 8]               0\n",
      "     BatchNorm2d-302            [-1, 640, 8, 8]           1,280\n",
      "            ReLU-303            [-1, 640, 8, 8]               0\n",
      "          Conv2d-304            [-1, 128, 8, 8]          20,480\n",
      "  CondensingConv-305            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-306            [-1, 128, 8, 8]             256\n",
      "            ReLU-307            [-1, 128, 8, 8]               0\n",
      "          Conv2d-308             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-309            [-1, 672, 8, 8]               0\n",
      "     BatchNorm2d-310            [-1, 672, 8, 8]           1,344\n",
      "            ReLU-311            [-1, 672, 8, 8]               0\n",
      "          Conv2d-312            [-1, 128, 8, 8]          21,504\n",
      "  CondensingConv-313            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-314            [-1, 128, 8, 8]             256\n",
      "            ReLU-315            [-1, 128, 8, 8]               0\n",
      "          Conv2d-316             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-317            [-1, 704, 8, 8]               0\n",
      "     BatchNorm2d-318            [-1, 704, 8, 8]           1,408\n",
      "            ReLU-319            [-1, 704, 8, 8]               0\n",
      "          Conv2d-320            [-1, 128, 8, 8]          22,528\n",
      "  CondensingConv-321            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-322            [-1, 128, 8, 8]             256\n",
      "            ReLU-323            [-1, 128, 8, 8]               0\n",
      "          Conv2d-324             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-325            [-1, 736, 8, 8]               0\n",
      "     BatchNorm2d-326            [-1, 736, 8, 8]           1,472\n",
      "            ReLU-327            [-1, 736, 8, 8]               0\n",
      "          Conv2d-328            [-1, 128, 8, 8]          23,552\n",
      "  CondensingConv-329            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-330            [-1, 128, 8, 8]             256\n",
      "            ReLU-331            [-1, 128, 8, 8]               0\n",
      "          Conv2d-332             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-333            [-1, 768, 8, 8]               0\n",
      "     BatchNorm2d-334            [-1, 768, 8, 8]           1,536\n",
      "            ReLU-335            [-1, 768, 8, 8]               0\n",
      "          Conv2d-336            [-1, 128, 8, 8]          24,576\n",
      "  CondensingConv-337            [-1, 128, 8, 8]               0\n",
      "     BatchNorm2d-338            [-1, 128, 8, 8]             256\n",
      "            ReLU-339            [-1, 128, 8, 8]               0\n",
      "          Conv2d-340             [-1, 32, 8, 8]           9,216\n",
      "     _DenseLayer-341            [-1, 800, 8, 8]               0\n",
      "     BatchNorm2d-342            [-1, 800, 8, 8]           1,600\n",
      "            ReLU-343            [-1, 800, 8, 8]               0\n",
      "       AvgPool2d-344            [-1, 800, 1, 1]               0\n",
      "          Linear-345                   [-1, 10]           4,010\n",
      "CondensingLinear-346                   [-1, 10]               0\n",
      "================================================================\n",
      "Total params: 516,202\n",
      "Trainable params: 516,202\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.01\n",
      "Forward/backward pass size (MB): 82.15\n",
      "Params size (MB): 1.97\n",
      "Estimated Total Size (MB): 84.13\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "from hdd.models.cnn.condensenet import convert_model\n",
    "\n",
    "convert_model(net)\n",
    "net = net.to(DEVICE)\n",
    "summary(net, (3, 32, 32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "eff167c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#Parameter: 516202\n",
      "Epoch: 1/15 Train Loss: 0.6438 Accuracy: 0.9385 Time: 29.19615  | Val Loss: 0.6686 Accuracy: 0.9314\n",
      "Epoch: 2/15 Train Loss: 0.6432 Accuracy: 0.9380 Time: 29.12915  | Val Loss: 0.6667 Accuracy: 0.9315\n",
      "Epoch: 3/15 Train Loss: 0.6434 Accuracy: 0.9386 Time: 29.01056  | Val Loss: 0.6673 Accuracy: 0.9332\n",
      "Epoch: 4/15 Train Loss: 0.6435 Accuracy: 0.9378 Time: 29.19215  | Val Loss: 0.6679 Accuracy: 0.9325\n",
      "Epoch: 5/15 Train Loss: 0.6448 Accuracy: 0.9387 Time: 28.78944  | Val Loss: 0.6654 Accuracy: 0.9346\n",
      "Epoch: 6/15 Train Loss: 0.6432 Accuracy: 0.9382 Time: 29.05139  | Val Loss: 0.6687 Accuracy: 0.9322\n",
      "Epoch: 7/15 Train Loss: 0.6479 Accuracy: 0.9369 Time: 29.09117  | Val Loss: 0.6674 Accuracy: 0.9304\n",
      "Epoch: 8/15 Train Loss: 0.6482 Accuracy: 0.9362 Time: 29.19816  | Val Loss: 0.6694 Accuracy: 0.9312\n",
      "Epoch: 9/15 Train Loss: 0.6432 Accuracy: 0.9387 Time: 29.14522  | Val Loss: 0.6682 Accuracy: 0.9334\n",
      "Epoch: 10/15 Train Loss: 0.6448 Accuracy: 0.9366 Time: 29.21176  | Val Loss: 0.6713 Accuracy: 0.9335\n",
      "Epoch: 11/15 Train Loss: 0.6450 Accuracy: 0.9371 Time: 28.79935  | Val Loss: 0.6691 Accuracy: 0.9314\n",
      "Epoch: 12/15 Train Loss: 0.6430 Accuracy: 0.9383 Time: 28.78810  | Val Loss: 0.6658 Accuracy: 0.9323\n",
      "Epoch: 13/15 Train Loss: 0.6440 Accuracy: 0.9383 Time: 28.83240  | Val Loss: 0.6651 Accuracy: 0.9318\n",
      "Epoch: 14/15 Train Loss: 0.6383 Accuracy: 0.9401 Time: 28.72094  | Val Loss: 0.6679 Accuracy: 0.9335\n",
      "Epoch: 15/15 Train Loss: 0.6378 Accuracy: 0.9411 Time: 29.14355  | Val Loss: 0.6666 Accuracy: 0.9340\n"
     ]
    }
   ],
   "source": [
    "from hdd.train.classification_utils import naive_train_classification_model\n",
    "\n",
    "print(f\"#Parameter: {count_trainable_parameter(net)}\")\n",
    "criteria = nn.CrossEntropyLoss(label_smoothing=0.1)\n",
    "optimizer = torch.optim.Adam(\n",
    "    net.parameters(),\n",
    "    lr=1e-4,\n",
    "    betas=(0.9, 0.99),\n",
    ")\n",
    "max_epochs = 15\n",
    "base_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
    "    optimizer, max_epochs, eta_min=1e-5\n",
    ")\n",
    "scheduler = GradualWarmupScheduler(\n",
    "    optimizer,\n",
    "    multiplier=1.0,\n",
    "    total_epoch=5,\n",
    "    after_scheduler=base_scheduler,\n",
    ")\n",
    "patch_4 = naive_train_classification_model(\n",
    "    net,\n",
    "    criteria,\n",
    "    max_epochs,\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    DEVICE,\n",
    "    optimizer,\n",
    "    scheduler,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d2a835a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-cu124",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
